/*
 有缺陷 不能压缩空格和换行 -> 更换文件读取方式
    - 生成频度表文件
    - 读取频度表文件
 结尾压缩不能够和好的处理
 最后不要忘记去除内存泄漏
*/

#include<iostream>
#include<unordered_map>
#include<vector>
#include<string>
#include<algorithm>
#include<fstream>
using namespace std;

//源文件
#define SOURCE_FILE "source.txt"
//频度文件
#define FREQUENCE_FILE "fre.txt"
//编码文件 用于解压缩
#define ENCODE_FILE "encode.txt"
//压缩文件
#define COMPRESS_FILE "compress.bin"
//解压缩文件
#define UMZIP_FILE "unzip.txt"
//树节点
typedef struct HuffmanNode{
    //构造函数 叶子
    HuffmanNode(char k,double w) : key(k), weight(w), left(nullptr), right(nullptr) {};
    //构造函数 非叶
    HuffmanNode(double w) : key('\0'), weight(w), left(nullptr), right(nullptr) {};
    //解码时使用 只需要字符的叶子节点
    HuffmanNode(char k) : key(k), weight(0), left(nullptr), right(nullptr) {};
    //解码时使用 什么都不需要的非叶子节点
    HuffmanNode() : key('\0'), weight(0), left(nullptr), right(nullptr) {};
    //字符 权重
    char key;
    double weight;
    HuffmanNode* left;
    HuffmanNode* right;
}HuffmanNode;

//节点堆结构 比较函数 小根堆
class ComHuffmanNode{
public:
    bool operator()(const HuffmanNode* n1,const HuffmanNode* n2)
    {
        return n1->weight > n2->weight;
    }
};


class HuffmanTree{
private:
    //存储树节点
    vector<HuffmanNode*> vechufffman;
    //Haffman树根节点
    HuffmanNode* rootNode;
    //编码表
    unordered_map<char,string> haffmanCode;
// 构造函数
public:
    HuffmanTree(){
        rootNode = nullptr;
    }
    ~HuffmanTree(){
        this->deleteTreeNode(this->rootNode);
        cout << "memory free" << endl;
    }
//关键函数
public:
    //生成频度文件
    void generateFrequenceFile(const string& filename);
    //读入频度文件
    void readData(const string& filename);
    //建造编码Huffman树 根据频度文件
    void buildEncedeTree();
    //建造解码Huffman树 根据编码文件
    void buildDecodeTree(const string& filename);
    //生成Huffman编码
    void generateCode(HuffmanNode* root, string& code);
    //加密文件
    void encodeFile(const string& filename);
    //解密文件
    void decodeFile(const string& filename);
    //释放树内存
    void deleteTreeNode(HuffmanNode* root);
//辅助函数
public:
    //打印节点数组
    void printVecHuff(const vector<HuffmanNode*>& VH);
    //打印编码
    void printCode(const unordered_map<char,string>& haffmanCode);
    //文件错误
    void fileError(const string& filename);
};

void HuffmanTree::generateFrequenceFile(const string& filename)
{
    cout << "begin" << endl;
    fstream file;
    file.open(filename,ios::in);
    //设置不跳过空格及换行
    file.unsetf(ios::skipws);
    if(!file.is_open()) this->fileError(filename);
    //统计每个字符出现的次数
    unordered_map<char,int> umap;
    char c;
    while(file >> c){
        printf("c:-%c-\n",c);
        //该字符第一次出现 记一次数
        if(umap.find(c) == umap.end()){
            umap.insert(pair<char,int>(c,1));
        }else{
            //计数+1
            umap.find(c)->second = umap.find(c)->second + 1;
        }
    }
    file.close();
    //将该次数统计写入文件
    file.open(FREQUENCE_FILE,ios::out | ios::trunc);
    if(!file.is_open()) this->fileError(FREQUENCE_FILE);
    for(auto mit = umap.begin(); mit != umap.end(); mit++){
        cout << "map[" << mit->first << "] = " << mit->second << endl;
        file << mit->first << " " << mit->second << endl;
    }
    file.close();
}

void HuffmanTree::readData(const string& filename)
{
    fstream file;
    file.open(filename,ios::in);
    file.unsetf(ios::skipws);
    if(!file.is_open()){
        cout << "file " << filename << " not exist" << endl;
        exit(-1); 
    }
    //读入叶子 字符+空格+权重+换行
    char key;
    char space;
    double weight;
    char nstr;
    while(file >> key >> space >> weight >> nstr){
        cout << key << space << weight << nstr;
        vechufffman.push_back(new HuffmanNode(key,weight));
    }
    file.close();
    cout << "read from file" << endl;
    this->printVecHuff(this->vechufffman);
}

void HuffmanTree::buildEncedeTree()
{
    //建堆
    make_heap(this->vechufffman.begin(),this->vechufffman.end(),ComHuffmanNode());
    //生成堆
    cout << "generate small heap:" << endl;
    this->printVecHuff(this->vechufffman);
    //造树
    while(this->vechufffman.size() > 1){
        //最小值
        HuffmanNode* left = this->vechufffman.front();
        pop_heap(this->vechufffman.begin(),this->vechufffman.end(),ComHuffmanNode());
        this->vechufffman.pop_back();
        //次小值
        HuffmanNode* right = this->vechufffman.front();
        pop_heap(this->vechufffman.begin(),this->vechufffman.end(),ComHuffmanNode());
        this->vechufffman.pop_back();
        //合并成新节点 非叶节点 权值为两者之和
        HuffmanNode* parent = new HuffmanNode(left->weight+right->weight);
        parent->left = left;
        parent->right = right;
        //新节点入堆
        this->vechufffman.push_back(parent);
        push_heap(this->vechufffman.begin(),this->vechufffman.end(),ComHuffmanNode());
        cout << "leftnode:" << left->key << "(" << left->weight << ") " 
            << " rightcode:" << right->key << "(" << right->weight << ") " 
            << " commbine to " << parent->weight << endl;
    }
    //最后的非叶节点由rootNode继承
    this->rootNode = vechufffman.front();
    //生成编码
    string code = "";
    this->generateCode(this->rootNode,code);
    //打印编码
    cout << "generate the haffman code:" << endl;
    this->printCode(this->haffmanCode);
    //将编码写入encode.txt
    fstream file;
    file.open(ENCODE_FILE,ios::out | ios::trunc);
    if(!file.is_open()) this->fileError(ENCODE_FILE);
    for(auto mit = haffmanCode.begin(); mit != haffmanCode.end(); mit++){
        file << mit->first << " " << mit->second << endl;
    }
    file.close();
}

void HuffmanTree::buildDecodeTree(const string& filename)
{
    fstream file;
    file.open(filename,ios::in);
    if(!file.is_open()) this->fileError(filename);
    file.unsetf(ios::skipws);
    //存储编码与字符的对应关系
    unordered_map<char,string> mcode;
    //字符 空格 换行
    char c,space,nc;
    string code;
    cout << "read encode.txt" << endl;
    while(file >> c >> space >> code >> nc){
        cout << c << space << code << nc;
        mcode.insert(pair<char,string>(c,code));
    }
    cout << "mcode as follow:" << endl;
    this->printCode(mcode);

    HuffmanNode* root = new HuffmanNode();
    HuffmanNode* tem = root;
    cout << "build tree" << endl;
    //建树
    for(auto mit = mcode.begin(); mit != mcode.end(); mit++){
        c = mit->first;
        cout << "c=" << c << endl;
        code = mit->second;
        for(auto cit = code.begin(); cit != code.end(); cit++){
            //左分支
            if((*cit) == '0'){
                if(tem->left == nullptr) tem->left = new HuffmanNode();
                tem = tem->left;
            }else{
                //右分支
                if(tem->right == nullptr) tem->right = new HuffmanNode();
                tem = tem->right;
            }
        }
        tem->key = c;
        cout << "build leaf:" << c << endl;
        tem = root;
    }
    this->rootNode = root;
}

void HuffmanTree::generateCode(HuffmanNode* root, string& code)
{
    //叶子节点
    if(root->left == nullptr){
        this->haffmanCode.insert(pair<char,string>(root->key,code));
        return;
    }
    code.push_back('0');
    generateCode(root->left,code);
    code.pop_back();
    code.push_back('1');
    generateCode(root->right,code);
    code.pop_back();
}

void HuffmanTree::encodeFile(const string& filename)
{
    //待加密文件
    fstream file_source;
    file_source.unsetf(ios::skipws);
    file_source.open(filename,ios::in);
    if(!file_source.is_open()) this->fileError(filename);
    //存储密文文件
    fstream file_encode;
    file_encode.open(COMPRESS_FILE,ios::out | ios::binary | ios::trunc);
    if(!file_encode.is_open()) this->fileError(COMPRESS_FILE);
    //压缩编码
    string compress_str = "";
    //读入源文件字符
    char c;
    //8个bit组成一个char
    unsigned char bin = 0x00;
    //计数
    int binCount = 0;
    while(file_source >> c){
        string code = this->haffmanCode.find(c)->second;
        compress_str += code;
        //将code转成bit串,写入压缩文件
        for(auto cit = code.begin(); cit != code.end(); cit++){
            if((*cit) == '1'){
                bin = (bin<<1) | 0x01;
            }else{
                bin = bin<<1;
            }
            binCount++;
            //写入 归零
            if(binCount == 8){
                cout << "binCount:" << binCount << endl;
                printf("write:%2x\n",bin);
                file_encode << bin;
                bin = 0x00;
                binCount = 0;
            }
        }
    }
    //bin余下未满8 单独处理
    cout << "binCount:" << binCount << endl;
    if(binCount != 0) bin = bin<<(8-binCount);
    printf("write:%2x\n",bin);
    file_encode << bin;

    file_encode.close();
    file_source.close();
    cout << "right code:";
    for(int i=0; i<compress_str.length(); i++){
        if(i%8 == 0){
            cout << endl;
        }
        cout << compress_str[i];
    }
}

void HuffmanTree::decodeFile(const string& filename)
{
    fstream file;
    file.open(filename,ios::in | ios::binary);
    if(!file.is_open()) this->fileError(filename);
    fstream file1;
    file1.open(UMZIP_FILE,ios::out | ios::trunc);
    if(!file1.is_open()) this->fileError(UMZIP_FILE);
    HuffmanNode * root = this->rootNode;
    //解压字符串
    string decode_str = "";
    char c;
    cout << "DECODE PROCESS:" << endl;
    while(file >> c){
        //从高位往低位取bit
        for(int i=7; i>=0; i--){
            char scode = (c >> i) & 0x01;
            if(scode == 0x00){
                cout << "0";
                root = root->left;
            }else{
                cout << "1";
                root = root->right;
            }
            //叶子 写入解压缩文件 恢复根节点 继续寻路
            if(root->left == nullptr){
                cout << " ==> " << root->key << endl;
                decode_str += root->key;
                file1 << root->key;
                root = this->rootNode;
            }
        }
    }
    cout << "\ndecode str:" << endl << decode_str << endl;
    file.close();
    file1.close();
}

void HuffmanTree::deleteTreeNode(HuffmanNode* root)
{
    if(root == nullptr) return;
    deleteTreeNode(root->left);
    deleteTreeNode(root->right);
    delete root;
    root = nullptr;
}

void HuffmanTree::printVecHuff(const vector<HuffmanNode*>& VH)
{
    for(auto vit = VH.begin(); vit != VH.end(); vit++){
        cout << (*vit)->key << "  " << (*vit)->weight << endl;
    }
}

void HuffmanTree::printCode(const unordered_map<char,string>& haffmanCode)
{
    for(auto mit = haffmanCode.begin(); mit != haffmanCode.end(); mit++){
        cout << mit->first << " : " << mit->second << endl;
    }
}

void HuffmanTree::fileError(const string& filename)
{
    cout << "file " << filename << " open error" << endl;
    exit(-1);
}

int main(){

    /**
     * 加密需要源文件source.txt即可
            HuffmanTree* huff = new HuffmanTree();
            huff->generateFrequenceFile("source.txt");
            huff->readData(FREQUENCE_FILE);
            huff->buildTree();
            huff->encodeFile(SOURCE_FILE);
     * 解密需要压缩文件compress.bin 以及 encode.txt 或者说 fre.txt(但是这个不可靠)
            根据encode.txt生成树
            解码
    */
    HuffmanTree* huff = new HuffmanTree();
    // huff->generateFrequenceFile(SOURCE_FILE);
    // huff->readData(FREQUENCE_FILE);
    // huff->buildEncedeTree();
    // huff->encodeFile(SOURCE_FILE);
    huff->buildDecodeTree(ENCODE_FILE);
    huff->decodeFile(COMPRESS_FILE);
    if(huff) delete huff;
    huff = nullptr;
    return 0;
}